docs/automation/ai_training.md (1)

25-28: Doc still references the old dialog UI.

The new flow opens a drawer with a “Model Training” tab. Leaving “dialog” here misleads users. Please update the wording to match the drawer-based UI.

-## Configuring the Training
-To start a training, click on the `AI Analysis` button in the toolbar and select `Train AI model` from the dropdown menu.
-This will open a dialog where you can configure and start your training job.
+## Configuring the Training
+To start training, click the `AI Analysis` button in the toolbar and choose `Train AI model`.
+This opens a drawer with a `Model Training` tab where you can configure and start your training job.

docs/automation/ai_training.md (1)

56-58: Fix the typo in the closing paragraph.

“Find and overview” should be “find an overview.”

-Once the training is finished, you can find and overview of all your trained model from the `Admin` >`AI Models` page in the navbar.
+Once the training is finished, you can find an overview of all your trained models on the `Admin` > `AI Models` page in the navbar.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

frontend/javascripts/viewer/view/ai_jobs/hooks/fetch_annotation_infos.ts (4)

65-79: Stop pushing to (userBoundingBoxes || []); mutate the actual array

Guard above ensures userBoundingBoxes is an array, so drop the fallback to avoid mutating a temporary. Also return the variable directly.

-      if (annotation.task?.boundingBox) {
-        const existingIds = (userBoundingBoxes || []).map(({ id }) => id);
-        const largestId = existingIds.length > 0 ? Math.max(...existingIds) : -1;
-        (userBoundingBoxes || []).push({
+      if (annotation.task?.boundingBox) {
+        const existingIds = userBoundingBoxes.map(({ id }) => id);
+        const largestId = existingIds.length > 0 ? Math.max(...existingIds) : -1;
+        userBoundingBoxes.push({
           name: "Task Bounding Box",
           boundingBox: Utils.computeBoundingBoxFromBoundingBoxObject(annotation.task.boundingBox),
           color: [0, 0, 0],
           isVisible: true,
           id: largestId + 1,
         });
       }
@@
-        volumeTracingMags: volumeTracingMags,
-        userBoundingBoxes: userBoundingBoxes || [],
+        volumeTracingMags,
+        userBoundingBoxes,

Also applies to: 80-90, 98-104

---

15-113: Extract helpers to reduce cognitive load

Split into small helpers: resolveInputsToAnnotationIds(inputs) and buildAnnotationBundle(annotationId). This improves readability and testability.

---

1-1: File placement nit: not a React hook

This module isn't a hook. Consider moving from hooks/ to utils/ or api/.

---

21-27: Robust URL parsing (handle query/fragment/trailing slash) and clearer var name

Current extraction breaks for URLs like ".../123?foo=bar" or with trailing "/". Decode and strip query/fragment. Also align naming with intent (annotationId).

Apply:

-    if (taskOrAnnotationIdOrUrl.includes("/")) {
-      const lastSegment = taskOrAnnotationIdOrUrl.split("/").at(-1);
-      if (lastSegment) {
-        annotationIdsForTraining.push(lastSegment);
-      }
-    } else {
+    if (taskOrAnnotationIdOrUrl.includes("/")) {
+      // Extract annotation id from URL-like input (strip trailing slashes, query, fragment)
+      let lastPathSegment = "";
+      try {
+        const u = new URL(taskOrAnnotationIdOrUrl);
+        lastPathSegment = u.pathname.replace(/\/+$/, "").split("/").at(-1) ?? "";
+      } catch {
+        // Not an absolute URL; treat as path-like string
+        lastPathSegment = taskOrAnnotationIdOrUrl.replace(/\/+$/, "").split("/").at(-1) ?? "";
+      }
+      const annotationId = lastPathSegment.split(/[?#]/, 1)[0];
+      if (annotationId) {
+        annotationIdsForTraining.push(decodeURIComponent(annotationId));
+      } else {
+        Toast.warning(`Could not extract annotation id from "${taskOrAnnotationIdOrUrl}"`);
+      }
+    } else {

frontend/javascripts/viewer/view/ai_jobs/utils.ts (1)

325-343: Wrap getIntersectingMagList call in ai_training_data_selector.tsx in try/catch

getIntersectingMagList invokes getSegmentationLayerByHumanReadableName, which throws if the layer isn’t found; handling this prevents uncaught errors in the component.

frontend/javascripts/viewer/view/ai_jobs/hooks/fetch_annotation_infos.ts (4)

33-34: Avoid magic string for annotation state

Replace "Finished" with a shared enum/constant if available to prevent drift.

Do we have a central AnnotationState/TaskState enum in types? If yes, import and use it here.

---

48-50: Deduplicate annotation IDs before fetching to avoid duplicate requests/UI entries

Use a Set to dedupe.

-  const newAnnotationsWithDatasets = await Promise.all(
-    annotationIdsForTraining.map(async (annotationId) => {
+  const uniqueAnnotationIdsForTraining = Array.from(new Set(annotationIdsForTraining));
+  const newAnnotationsWithDatasets = await Promise.all(
+    uniqueAnnotationIdsForTraining.map(async (annotationId) => {

---

53-60: Use enum instead of string literal for layer type + parallelize dataset and tracing fetches

• Prefer AnnotationLayerEnum.Volume for consistency (you already use the enum for Skeleton).
• Fetch dataset and tracings concurrently to reduce latency.

-      const dataset = await getDataset(annotation.datasetId);
-
-      const volumeServerTracings: ServerVolumeTracing[] = await Promise.all(
-        annotation.annotationLayers
-          .filter((layer) => layer.typ === "Volume")
-          .map(
-            (layer) =>
-              getTracingForAnnotationType(annotation, layer) as Promise<ServerVolumeTracing>,
-          ),
-      );
+      const [dataset, volumeServerTracings] = await Promise.all([
+        getDataset(annotation.datasetId),
+        Promise.all(
+          annotation.annotationLayers
+            .filter((layer) => layer.typ === AnnotationLayerEnum.Volume)
+            .map(
+              (layer) =>
+                getTracingForAnnotationType(annotation, layer) as Promise<ServerVolumeTracing>,
+            ),
+        ),
+      ]);

---

19-20: De-dup unfinished tasks to avoid repeated names in the toast

Use Set to collect, then spread on display.

-  const unfinishedTasks: string[] = [];
+  const unfinishedTasks = new Set<string>();
@@
-          unfinishedTasks.push(taskOrAnnotationIdOrUrl);
+          unfinishedTasks.add(taskOrAnnotationIdOrUrl);
@@
-  if (unfinishedTasks.length > 0) {
-    Toast.warning(
-      `The following tasks have no finished annotations: ${unfinishedTasks.join(", ")}`,
-    );
+  if (unfinishedTasks.size > 0) {
+    Toast.warning(
+      `The following tasks have no finished annotations: ${[...unfinishedTasks].join(", ")}`,
+    );
   }

Also applies to: 37-38, 107-111

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

frontend/javascripts/viewer/view/action_bar_view.tsx (1)

298-334: Critical issue from previous review remains unfixed: onClick handlers on menu items are ignored in antd v5.

The onClick handlers on individual menu items (lines 301, 308, 315) are ignored by antd v5's Dropdown component. The AI Analysis dropdown buttons will not work. This was flagged as a critical issue in the previous review and the fix was not applied.

Apply the previously suggested fix by moving onClick logic to menu.onClick:

 const isSuperUser = this.props.activeUser?.isSuperUser === true;
 const menuItems = [
-  {
-    key: "open_ai_inference_button",
-    onClick: () => Store.dispatch(setAIJobDrawerStateAction("open_ai_inference")),
-    label: "Run AI model",
-  },
-  ...(isSuperUser
-    ? [
-        {
-          key: "open_ai_training_button",
-          onClick: () => Store.dispatch(setAIJobDrawerStateAction("open_ai_training")),
-          label: "Train new AI model",
-        },
-      ]
-    : []),
-  {
-    key: "open_ai_alignment_button",
-    onClick: () => Store.dispatch(setAIJobDrawerStateAction("open_ai_alignment")),
-    label: "Run AI Alignment",
-  },
+  { key: "open_ai_inference_button", label: "Run AI model" },
+  ...(isSuperUser ? [{ key: "open_ai_training_button", label: "Train new AI model" }] : []),
+  { key: "open_ai_alignment_button", label: "Run AI Alignment" },
 ];

 return (
   <div>
     <Dropdown
       key="ai-job-dropdown"
       menu={{
         items: menuItems,
+        onClick: ({ key }) => {
+          switch (key) {
+            case "open_ai_inference_button":
+              Store.dispatch(setAIJobDrawerStateAction("open_ai_inference"));
+              break;
+            case "open_ai_training_button":
+              Store.dispatch(setAIJobDrawerStateAction("open_ai_training"));
+              break;
+            case "open_ai_alignment_button":
+              Store.dispatch(setAIJobDrawerStateAction("open_ai_alignment"));
+              break;
+          }
+        },
       }}
       disabled={disabled}
     >
       <Button disabled={disabled} icon={<i className="fas fa-magic" />} title={tooltipText}>
         AI Analysis
       </Button>
     </Dropdown>
   </div>
 );

Based on learnings

frontend/javascripts/viewer/view/ai_jobs/alignment/ai_alignment_job_context.tsx (3)

41-44: Add null checks for colorLayer to prevent runtime errors.

If a dataset has no color layers, colorLayers[0] will be undefined, causing getBoundingBoxesForLayers([colorLayer])[0] to fail and colorLayer.name (line 61) to throw. This issue was previously flagged but remains unresolved.

Apply this diff to add guards:

   const dataset = useWkSelector((state) => state.dataset);
   const annotationId = useWkSelector((state) => state.annotation.annotationId);
   const colorLayers = getColorLayers(dataset);
-  const colorLayer = colorLayers[0];
+  const colorLayer = colorLayers[0] ?? null;

-  const selectedBoundingBox = getBoundingBoxesForLayers([colorLayer])[0];
+  const selectedBoundingBox = colorLayer ? getBoundingBoxesForLayers([colorLayer])[0] : null;

---

52-55: Update validation to include colorLayer check.

The validation doesn't check whether colorLayer exists, but handleStartAnalysis requires it. This will allow the Start button to be enabled even when the job will fail.

Apply this diff:

   const areParametersValid = useMemo(
-    () => Boolean(selectedTask && selectedJobType && newDatasetName),
-    [selectedTask, selectedJobType, newDatasetName],
+    () => Boolean(selectedTask && selectedJobType && newDatasetName && colorLayer),
+    [selectedTask, selectedJobType, newDatasetName, colorLayer],
   );

---

57-71: Add early return guard for invalid parameters.

The handler should verify parameters before attempting the API call. Without this guard, the function will crash on colorLayer.name when colorLayer is null.

Apply this diff:

   const handleStartAnalysis = useCallback(async () => {
+    if (!areParametersValid || !colorLayer) {
+      Toast.error("Please ensure all required parameters are set and a color layer exists.");
+      return;
+    }
     try {
       await startAlignSectionsJob(
         dataset.id,
-        colorLayer.name,
+        colorLayer!.name,
         newDatasetName,
         shouldUseManualMatches ? annotationId : undefined,
       );
       Toast.success("Alignment started successfully!");
       dispatch(setAIJobDrawerStateAction("invisible"));
     } catch (error) {
       console.error(error);
       Toast.error("Failed to start alignment.");
     }
-  }, [dataset.id, dispatch, colorLayer.name, newDatasetName, annotationId, shouldUseManualMatches]);
+  }, [dataset.id, dispatch, colorLayer, newDatasetName, annotationId, shouldUseManualMatches, areParametersValid]);

frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_analysis_settings.tsx (3)

149-158: Type‑narrow before mag/size validation.

Guard to inference job types to avoid invalid calls.

-                if (value && selectedLayer && selectedJobType) {
+                if (
+                  value &&
+                  selectedLayer &&
+                  (selectedJobType === APIJobType.INFER_MITOCHONDRIA ||
+                    selectedJobType === APIJobType.INFER_NEURONS ||
+                    selectedJobType === APIJobType.INFER_NUCLEI ||
+                    selectedJobType === APIJobType.INFER_INSTANCES)
+                ) {

---

61-65: Fix layer change handler: Select returns a string name.

-    if ("selectedLayer" in changedValues) {
-      setSelectedLayer(
-        colorLayers.find((l) => l.name === changedValues.selectedLayer.name) as APIDataLayer,
-      );
-    }
+    if ("selectedLayer" in changedValues) {
+      const name = changedValues.selectedLayer as string;
+      setSelectedLayer(colorLayers.find((l) => l.name === name) as APIDataLayer);
+    }

---

127-139: Make uint24 validation effective for Select.

The field stores a layer name (string). Validate against the resolved selectedLayer.

             <Form.Item
               name="selectedLayer"
               label="Image Data Layer"
               rules={[
                 { required: true, message: "Please select an image data layer" },
-                colorLayerMustNotBeUint24Rule,
+                {
+                  validator: () => {
+                    if (selectedLayer && (selectedLayer as any).elementClass === "uint24") {
+                      return Promise.reject(
+                        new Error(
+                          "The selected layer of type uint24 is not supported. Please select a different one.",
+                        ),
+                      );
+                    }
+                    return Promise.resolve();
+                  },
+                },
               ]}
             >

docs/automation/ai_training.md (1)

26-28: Update “dialog” to “drawer” and reference the tab.

The new UX is a drawer with a “Model Training” tab; adjust copy accordingly.

-## Configuring the Training
-To start a training, click on the `AI Analysis` button in the toolbar and select `Train AI model` from the dropdown menu.
-This will open a dialog where you can configure and start your training job.
+## Configuring the Training
+To start training, click the `AI Analysis` button in the toolbar and select `Train AI model`.
+This opens a drawer with a `Model Training` tab where you can configure and start your training job.

frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx (1)

54-65: Fix category→jobType mapping; include NUCLEI and MITOCHONDRIA.

EM_NUCLEI maps to INFER_NUCLEI (not INSTANCES). Also add EM_MITOCHONDRIA → INFER_MITOCHONDRIA to avoid wrong pipelines/costs and errors on selection.

-const mapCategoryToJobType = (
-  category: APIAiModelCategory,
-): APIJobType.INFER_NEURONS | APIJobType.INFER_INSTANCES => {
+const mapCategoryToJobType = (
+  category: APIAiModelCategory,
+): AllowedJobType => {
   switch (category) {
     case APIAiModelCategory.EM_NEURONS:
       return APIJobType.INFER_NEURONS;
     case APIAiModelCategory.EM_NUCLEI:
-      return APIJobType.INFER_INSTANCES;
+      return APIJobType.INFER_NUCLEI;
+    case APIAiModelCategory.EM_MITOCHONDRIA:
+      return APIJobType.INFER_MITOCHONDRIA;
     default:
       throw new Error(`Unsupported category: ${category}`);
   }
 };

Also update types where used:

Type alias (near Lines 13-24):

type AllowedJobType =
  | APIJobType.INFER_NEURONS
  | APIJobType.INFER_NUCLEI
  | APIJobType.INFER_MITOCHONDRIA
  | APIJobType.INFER_INSTANCES;

type PretrainedModel = { /* ... */ jobType: AllowedJobType; /* ... */ };

onSelectModel (Lines 84-88):

let jobType: AllowedJobType;

frontend/javascripts/viewer/view/ai_jobs/train_ai_model/ai_training_job_context.tsx (3)

11-11: Remove lodash/every import; use native checks.

Drop the dependency and avoid truthiness pitfalls.

-import every from "lodash/every";

---

78-82: Avoid non-null assertion in queryFn.

Read id from queryKey to drop the unsafe "!".

-  const { data: initialFullAnnotation } = useQuery({
-    queryKey: ["initialAnnotation", currentAnnotation.annotationId],
-    queryFn: () => getUnversionedAnnotationInformation(currentAnnotation.annotationId!),
-    enabled: !!currentAnnotation.annotationId,
-  });
+  const { data: initialFullAnnotation } = useQuery({
+    queryKey: ["initialAnnotation", currentAnnotation.annotationId],
+    queryFn: ({ queryKey }) => getUnversionedAnnotationInformation(queryKey[1] as string),
+    enabled: Boolean(currentAnnotation.annotationId),
+  });

---

124-131: Replace fragile validity check; accept enum 0 and validate Vector3.

Current check can misclassify valid values and blocks TS narrowing. Use explicit checks and a type guard.

-  const areParametersValid = useMemo(() => {
-    const areSelectionsValid = every(
-      selectedAnnotations,
-      (s) => s.imageDataLayer && s.groundTruthLayer && s.magnification,
-    );
-
-    return every([modelName, selectedJobType, areSelectionsValid, selectedAnnotations.length > 0]);
-  }, [modelName, selectedJobType, selectedAnnotations]);
+  const areParametersValid = useMemo(() => {
+    const hasAtLeastOne = selectedAnnotations.length > 0;
+    const areSelectionsValid = selectedAnnotations.every(isCompleteSelection);
+    return modelName.trim().length > 0 && selectedJobType != null && hasAtLeastOne && areSelectionsValid;
+  }, [modelName, selectedJobType, selectedAnnotations]);

Add near the interfaces:

type CompleteSelection = AiTrainingAnnotationSelection &
  Required<Pick<AiTrainingAnnotationSelection, "imageDataLayer" | "groundTruthLayer" | "magnification">>;

const isVector3 = (v: unknown): v is Vector3 =>
  Array.isArray(v) && v.length === 3 && v.every((n) => typeof n === "number" && Number.isFinite(n));

const isCompleteSelection = (s: AiTrainingAnnotationSelection): s is CompleteSelection =>
  typeof s.imageDataLayer === "string" &&
  s.imageDataLayer.length > 0 &&
  typeof s.groundTruthLayer === "string" &&
  s.groundTruthLayer.length > 0 &&
  isVector3(s.magnification);

Based on learnings

frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_model_selector.tsx (1)

224-228: Label custom model categories precisely.

Show NUCLEI and MITOCHONDRIA tags instead of generic INSTANCES.

-                    <Tag>
-                      {(item as AiModel).category === APIAiModelCategory.EM_NEURONS
-                        ? "NEURONS"
-                        : "INSTANCES"}
-                    </Tag>
+                    <Tag>
+                      {(item as AiModel).category === APIAiModelCategory.EM_NEURONS
+                        ? "NEURONS"
+                        : (item as AiModel).category === APIAiModelCategory.EM_NUCLEI
+                        ? "NUCLEI"
+                        : (item as AiModel).category === APIAiModelCategory.EM_MITOCHONDRIA
+                        ? "MITOCHONDRIA"
+                        : "INSTANCES"}
+                    </Tag>

frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_analysis_settings.tsx (1)

191-197: Avoid defaultValue on a controlled Form field.

Remove defaultValue to prevent controlled/uncontrolled mismatch.

-                    <InputNumber
+                    <InputNumber
                       min={0.1}
                       suffix="nm"
                       style={{ width: "100%" }}
-                      defaultValue={1000}
                       disabled={!isInstanceModel}
                     />

frontend/javascripts/viewer/view/ai_jobs/train_ai_model/ai_training_job_context.tsx (2)

148-152: Typo: rename commonJobArgmuments → commonJobArguments.

Small clarity tweak.

-    const commonJobArgmuments = {
+    const commonJobArguments = {
       trainingAnnotations: trainingAnnotations,
       name: modelName,
       comment: comments,
     };
@@
-          ...commonJobArgmuments,
+          ...commonJobArguments,
@@
-          ...commonJobArgmuments,
+          ...commonJobArguments,

Also applies to: 159-159, 165-165

---

183-188: Expose setters with the declared narrower signatures.

Wrap useState dispatchers to match AiTrainingJobContextType and avoid passing the broader SetStateAction signature.

   const value = {
     selectedJobType,
     selectedTask,
-    setSelectedJobType,
-    setSelectedTask,
+    setSelectedJobType: setSelectedJobTypeValue,
+    setSelectedTask: setSelectedTaskValue,
     handleStartAnalysis,

Add near the other callbacks:

const setSelectedJobTypeValue = useCallback((jobType: APIJobType) => setSelectedJobType(jobType), []);
const setSelectedTaskValue = useCallback((task: AiTrainingTask) => setSelectedTask(task), []);

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: knollengewaechs
PR: scalableminds/webknossos#8796
File: frontend/javascripts/viewer/view/action-bar/ai_job_modals/forms/neuron_segmentation_form.tsx:68-77
Timestamp: 2025-08-27T11:32:54.033Z
Learning: In the webknossos frontend AI job forms, `datasetConfiguration.layers[colorLayer.name]` is assumed to always be non-null when accessing layer properties like `isInverted`. The selected color layer is guaranteed to exist in the dataset configuration.

frontend/javascripts/viewer/view/ai_jobs/hooks/fetch_annotation_infos.ts (1)

20-21: Clarify support for "task URLs" in documentation.

The function comment claims to parse "task IDs, annotation IDs, or annotation URLs," but past review discussions indicate that task URLs don't exist in the system. The current implementation only handles annotation URLs (lines 27-31) by extracting the annotation ID from the path. Consider updating the JSDoc and parameter name to accurately reflect this limitation.

frontend/javascripts/viewer/view/ai_jobs/hooks/fetch_annotation_infos.ts (2)

27-31: Strengthen URL parsing to handle edge cases.

The current split-by-slash approach may fail if the URL contains trailing slashes, query parameters, or fragment identifiers. A URL ending in /# or including ?param=value could produce an incorrect or empty annotation ID.

Consider using the URL API for more robust parsing:

     if (taskOrAnnotationIdOrUrl.includes("/")) {
-      const annotationId = taskOrAnnotationIdOrUrl.split("/").at(-1);
+      try {
+        const url = new URL(taskOrAnnotationIdOrUrl, window.location.origin);
+        const annotationId = url.pathname.split("/").filter(Boolean).at(-1);
+      } catch {
+        // If URL parsing fails, fall back to simple split
+        const annotationId = taskOrAnnotationIdOrUrl.split("/").filter(Boolean).at(-1);
+      }
       if (annotationId) {
         annotationIdsForTraining.push(annotationId);
       }

---

86-120: LGTM! Consider clarifying the comment about the original bug.

The bounding box extraction logic is solid: it correctly handles empty volume tracings via optional chaining, falls back to skeleton layers when needed, throws an appropriate error if neither exists, and properly appends the task bounding box.

However, the comment on line 93 references "the original code had a bug" which might confuse future readers since the bug is no longer present.

Consider updating the comment to be more forward-looking:

-    // The original code had a bug here (!userBoundingBoxes) which was always false for an empty array.
+    // Check if the array is empty (not just falsy)
     const skeletonLayer = annotation.annotationLayers.find(

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

frontend/javascripts/viewer/view/ai_jobs/train_ai_model/ai_training_job_context.tsx (2)

11-11: Replace lodash every with explicit validation.

The lodash every call is fragile because it relies on truthiness, which can incorrectly reject valid falsy values (e.g., 0 for an enum). Additionally, every([]) returns true, though you mitigate this separately. Use explicit boolean checks for clarity and correctness.

Apply this diff:

-import every from "lodash/every";
 import type React from "react";

   const areParametersValid = useMemo(() => {
-    const areSelectionsValid = every(
-      selectedAnnotations,
-      (s) => s.imageDataLayer && s.groundTruthLayer && s.magnification,
-    );
-
-    return every([modelName, selectedJobType, areSelectionsValid, selectedAnnotations.length > 0]);
+    const areSelectionsValid =
+      selectedAnnotations.length > 0 &&
+      selectedAnnotations.every(
+        (s) =>
+          Boolean(s.imageDataLayer) &&
+          Boolean(s.groundTruthLayer) &&
+          Array.isArray(s.magnification) &&
+          s.magnification.length === 3,
+      );
+    return modelName.trim().length > 0 && selectedJobType != null && areSelectionsValid;
   }, [modelName, selectedJobType, selectedAnnotations]);

Also applies to: 124-131

---

139-146: Filter for complete selections before mapping.

TypeScript fails because optional fields can be undefined, but AiModelTrainingAnnotationSpecification requires all fields to be defined. Filter the selections first using a type guard.

Add a type guard before the mapping:

+    // Type guard to ensure all required fields are present
+    const isCompleteSelection = (
+      selection: AiTrainingAnnotationSelection,
+    ): selection is Required<AiTrainingAnnotationSelection> =>
+      Boolean(selection.imageDataLayer) &&
+      Boolean(selection.groundTruthLayer) &&
+      Boolean(selection.magnification);
+
+    const completeSelections = selectedAnnotations.filter(isCompleteSelection);
+    if (completeSelections.length !== selectedAnnotations.length) {
+      Toast.error("Some annotations are missing required fields.");
+      return;
+    }
+
-    const trainingAnnotations: AiModelTrainingAnnotationSpecification[] = selectedAnnotations.map(
+    const trainingAnnotations: AiModelTrainingAnnotationSpecification[] = completeSelections.map(
       (selection) => ({

frontend/javascripts/viewer/view/ai_jobs/train_ai_model/ai_training_model_selector.tsx (1)

24-24: Fix grammar: add missing preposition.

The phrase "Optimized dense neuronal tissue" is missing "for" before "dense."

Apply this diff:

-      "Train a new AI model for EM neuron segmentation based on the annotations in this dataset. Optimized dense neuronal tissue from SEM, FIB-SEM, SBEM, Multi-SEM microscopes.",
+      "Train a new AI model for EM neuron segmentation based on the annotations in this dataset. Optimized for dense neuronal tissue from SEM, FIB-SEM, SBEM, Multi-SEM microscopes.",

frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_analysis_settings.tsx (3)

61-65: Critical: Fix the selectedLayer change handler.

The Select component returns the layer name as a string (because options maps value: l.name), but line 63 incorrectly accesses changedValues.selectedLayer.name. This will cause selectedLayer to remain undefined.

Apply this diff:

     if ("selectedLayer" in changedValues) {
       setSelectedLayer(
-        colorLayers.find((l) => l.name === changedValues.selectedLayer.name) as APIDataLayer,
+        colorLayers.find((l) => l.name === changedValues.selectedLayer) as APIDataLayer,
       );
     }

---

127-139: Critical: uint24 validation is ineffective.

The field stores a string (layer name), but colorLayerMustNotBeUint24Rule expects an APIDataLayer object. The validator will never reject uint24 layers because value.elementClass is undefined on strings.

Apply this diff to validate against the resolved layer:

             <Form.Item
               name="selectedLayer"
               label="Image Data Layer"
               rules={[
                 { required: true, message: "Please select an image data layer" },
-                colorLayerMustNotBeUint24Rule,
+                {
+                  validator: () => {
+                    if (selectedLayer && selectedLayer.elementClass === "uint24") {
+                      return Promise.reject(
+                        new Error(
+                          "The selected layer of type uint24 is not supported. Please select a different one.",
+                        ),
+                      );
+                    }
+                    return Promise.resolve();
+                  },
+                },
               ]}
             >

---

149-169: Critical: Add type guard before calling inference-only functions.

Both getBestFittingMagComparedToTrainingDS and isDatasetOrBoundingBoxTooSmall accept only inference job types (INFER_MITOCHONDRIA | INFER_NEURONS | INFER_NUCLEI | INFER_INSTANCES). Calling them with a training or alignment job type will cause type errors.

Add a type guard at line 150:

             {
               validator: (_, value: UserBoundingBox) => {
                 if (value && selectedLayer && selectedJobType) {
+                  // Only validate for inference jobs
+                  const inferenceJobTypes = ["infer_mitochondria", "infer_neurons", "infer_nuclei", "infer_instances"];
+                  if (!inferenceJobTypes.includes(selectedJobType)) {
+                    return Promise.resolve();
+                  }
+
                   const boundingBox = computeArrayFromBoundingBox(value.boundingBox);
 
                   const mag = getBestFittingMagComparedToTrainingDS(
                     selectedLayer,
                     dataset.dataSource.scale,
                     selectedJobType,
                   );

Alternatively, if selectedJobType has a discriminated union type, use a type predicate or check a discriminant property.

docs/automation/ai_training.md (2)

26-27: Update UI copy: modal → drawer, include tab name.

The UI now uses a drawer with a “Model Training” tab, not a dialog.

-To start a training, click on the `AI Analysis` button in the toolbar and select `Train AI model` from the dropdown menu.
-This will open a dialog where you can configure and start your training job.
+To start training, click the `AI Analysis` button in the toolbar and choose `Train AI model`.
+This opens a drawer with a `Model Training` tab where you can configure and start your training job.

---

7-8: Tighten style and fix “on‑premise” → “on‑premises”.

-    AI Model Training is only available on [webknossos.org](https://webknossos.org) at the moment. 
-    If you want to set up on-premise automated analysis at your institute/workplace, then [please contact sales](mailto:[email protected]). 
+    AI Model Training is currently available on [webknossos.org](https://webknossos.org).
+    If you want to set up on‑premises automated analysis at your institute/workplace, [please contact sales](mailto:[email protected]).

frontend/javascripts/viewer/view/ai_jobs/train_ai_model/ai_training_job_context.tsx (1)

80-80: Avoid non-null assertion; use type guard.

The ! is safe here because the query is enabled only when annotationId is truthy, but non-null assertions can mask refactoring risks and are a code smell.

Apply this diff to use optional chaining instead:

-    queryFn: () => getUnversionedAnnotationInformation(currentAnnotation.annotationId!),
+    queryFn: () =>
+      currentAnnotation.annotationId
+        ? getUnversionedAnnotationInformation(currentAnnotation.annotationId)
+        : Promise.reject(new Error("No annotation ID")),
     enabled: !!currentAnnotation.annotationId,

docs/automation/ai_training.md (2)

23-23: Grammar/clarity: labeling paragraph.

-3.  **Label segments:** Within your bounding boxes, label the segmentation of your structures of interest. Use the volume annotation tool to manually annotate structures. This will be your ground truth data. For neurons, we recommend to densely label each structure with a unique ID. For instances segmentations you only need to label the structures you want to train on, e.g. nuclei, mitochondria, soma, vesicles, etc. 
+3.  **Label segments:** Within your bounding boxes, label the structures of interest using the volume annotation tool. This will be your ground truth data. For neurons, we recommend densely labeling each structure with a unique ID. For instance segmentations, you only need to label the structures you want to train on (e.g., nuclei, mitochondria, soma, vesicles).

---

48-48: Units/examples formatting.

-*   **Max Distance (nm):** (Only for EM Instances Model) The maximum cross-section length ("diameter") for each identified object in nanometers e.g. Nuclei: 1000nm, Vesicles: 80nm.
+*   **Max Distance (nm):** (Only for EM Instances Model) The maximum cross‑section length (“diameter”) for each identified object in nanometers (e.g., nuclei: 1,000 nm; vesicles: 80 nm).

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

frontend/javascripts/viewer/view/ai_jobs/alignment/ai_alignment_job_context.tsx (1)

37-41: Critical issue from previous review remains unresolved: guard against missing color layers.

The code still assumes colorLayers[0] exists. If the dataset has no color layers, colorLayer is undefined and line 40 crashes when deriving selectedBoundingBox. Lines 48-51 validation and lines 53-67 handleStartAnalysis also fail to check colorLayer existence before using colorLayer.name.

Apply this diff to add guards:

   const dataset = useWkSelector((state) => state.dataset);
   const annotationId = useWkSelector((state) => state.annotation.annotationId);
   const colorLayers = getColorLayers(dataset);
-  const colorLayer = colorLayers[0];
+  const colorLayer = colorLayers[0] ?? null;
 
-  const selectedBoundingBox = getBoundingBoxesForLayers([colorLayer])[0];
+  const selectedBoundingBox = colorLayer ? getBoundingBoxesForLayers([colorLayer])[0] : null;

Update validation to include colorLayer check:

   const areParametersValid = useMemo(
-    () => Boolean(selectedTask?.jobType && newDatasetName),
-    [selectedTask, newDatasetName],
+    () => Boolean(selectedTask?.jobType && newDatasetName && colorLayer),
+    [selectedTask, newDatasetName, colorLayer],
   );

Short-circuit start when colorLayer is missing:

   const handleStartAnalysis = useCallback(async () => {
+    if (!colorLayer) {
+      Toast.error("No color layer available. Please ensure the dataset has a color layer.");
+      return;
+    }
     try {
       await startAlignSectionsJob(
         dataset.id,
         colorLayer.name,

docs/automation/ai_segmentation.md (1)

7-7: Unresolved from previous review: terminology mismatch with UI.

The text still says "dialog" but the UI uses a "drawer". Update to match the implementation.

Apply this diff:

-You can launch the AI analysis dialog using the `AI Analysis` button in the toolbar at the top. This will open a dropdown menu with three options:
+You can launch the AI analysis using the `AI Analysis` button in the toolbar at the top. This will open a dropdown menu with three options. Selecting an option opens the AI Jobs drawer:

frontend/javascripts/viewer/view/ai_jobs/credit_information.tsx (2)

65-65: Remove unprofessional comment.

The inline comment "This is a shitty way to do it" should be removed or reworded professionally. If the approach is temporary, use a standard TODO or FIXME marker.

Apply this diff:

-  // sum all training volumes into a single bounding box
-  // This is a shitty way to do it, but it works for now.
+  // Sum all training volumes into a single synthetic bounding box for cost estimation
   const totalVolume = selectedAnnotations.reduce(

---

64-86: Synthetic cubic bounding box obscures actual training data distribution.

The code sums volumes from all training annotations and creates a synthetic cube with side length Math.cbrt(totalVolume). This approach loses information about the actual spatial distribution and shape of the training data. Consider whether the cost API can accept multiple bounding boxes or whether documentation should clarify this is an approximation.

If the API supports it, pass multiple bounding boxes and sum costs. Otherwise, add a comment explaining the cubic approximation:

+  // Note: Training cost is estimated using a synthetic cubic bounding box with
+  // volume equal to the sum of all training annotation volumes. This approximates
+  // the actual spatial distribution for cost calculation purposes.
   const totalVolume = selectedAnnotations.reduce(

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Learnt from: knollengewaechs
PR: scalableminds/webknossos#8796
File: frontend/javascripts/viewer/view/action-bar/ai_job_modals/forms/neuron_segmentation_form.tsx:68-77
Timestamp: 2025-08-27T11:32:54.033Z
Learning: In the webknossos frontend AI job forms, `datasetConfiguration.layers[colorLayer.name]` is assumed to always be non-null when accessing layer properties like `isInverted`. The selected color layer is guaranteed to exist in the dataset configuration.

frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_image_segmentation_job_context.tsx (1)

164-191: Fix brittle custom vs. pre-trained model detection.

The check "trainingJob" in selectedModel returns true for all AiModel objects (since trainingJob is always present as a property), including pre-trained models where trainingJob: null. This causes pre-trained models to incorrectly route to the custom inference path.

Apply this diff to check for non-null trainingJob:

-      if ("trainingJob" in selectedModel!) {
+      if (selectedModel!.trainingJob != null) {
         // Custom models

frontend/javascripts/viewer/view/ai_jobs/run_ai_model/ai_image_segmentation_job_context.tsx (2)

168-168: Remove unnecessary type assertion.

The id property is already typed as string in the AiModel interface, so the as string assertion is redundant.

Apply this diff:

-          aiModelId: selectedModel!.id as string,
+          aiModelId: selectedModel!.id,

---

121-133: Consider a type guard function for better type safety.

While the areParametersValid check with early return prevents runtime issues, TypeScript cannot infer type narrowing from lodash.every, requiring ! operators throughout (lines 137, 161, 164, 168, 171, 198, 210, 216, 218). A custom type predicate would eliminate these assertions.

Example:

function assertRequiredParameters(
  model: AiModel | Partial<AiModel> | null,
  jobType: APIJobType | null,
  bbox: UserBoundingBox | null,
  datasetName: string,
  layer: APIDataLayer | null,
): asserts model is AiModel | Partial<AiModel> {
  if (!model || !jobType || !bbox || !datasetName || !layer) {
    throw new Error("Missing required parameters");
  }
}

// Then in handleStartAnalysis:
try {
  assertRequiredParameters(selectedModel, selectedJobType, selectedBoundingBox, newDatasetName, selectedLayer);
  // TypeScript now knows these are non-null
  const boundingBox = computeArrayFromBoundingBox(selectedBoundingBox.boundingBox);
  // ...
}

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

frontend/javascripts/viewer/view/ai_jobs/credit_information.tsx (2)

43-43: Consider renaming for consistency.

The variable selectJobType could be renamed to selectedJobType to match the naming pattern used in other wrappers and the prop it's passed to (line 48).

Apply this diff:

-  const selectJobType = selectedTask?.jobType ?? null;
+  const selectedJobType = selectedTask?.jobType ?? null;

And update line 48:

-      selectedJobType={selectJobType}
+      selectedJobType={selectedJobType}

---

143-155: Use a stable key for the bounding box in the query.

The query key currently uses the object reference selectedBoundingBox?.boundingBox, which can cause unnecessary refetches when the object reference changes even if the content remains identical. Use a stable, serializable representation instead.

Apply this diff:

   const { data: jobCreditCostInfo, isFetching } = useQuery<JobCreditCostInfo>({
     queryKey: [
       "jobCreditCost",
       selectedJobType ?? "no-type",
-      selectedBoundingBox?.boundingBox ?? "no-bb",
+      selectedBoundingBox
+        ? computeArrayFromBoundingBox(selectedBoundingBox.boundingBox).join(",")
+        : "no-bb",
     ],
     queryFn: async () =>
       await getJobCreditCost(
         selectedJobType!,
         computeArrayFromBoundingBox(selectedBoundingBox!.boundingBox),
       ),
     enabled: Boolean(selectedBoundingBox && selectedJobType),
   });

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro